{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7228a58b-4f81-4f5d-ac6c-d9439b3f4447",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: HF_ENDPOINT=https://hf-mirror.com\n"
     ]
    }
   ],
   "source": [
    "%env HF_ENDPOINT=https://hf-mirror.com"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9ff5004e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/lvyufeng/miniconda3/envs/mindspore/lib/python3.9/site-packages/numpy/core/getlimits.py:549: UserWarning: The value of the smallest subnormal for <class 'numpy.float64'> type is zero.\n",
      "  setattr(self, word, getattr(machar, word).flat[0])\n",
      "/home/lvyufeng/miniconda3/envs/mindspore/lib/python3.9/site-packages/numpy/core/getlimits.py:89: UserWarning: The value of the smallest subnormal for <class 'numpy.float64'> type is zero.\n",
      "  return self._float_to_str(self.smallest_subnormal)\n",
      "/home/lvyufeng/miniconda3/envs/mindspore/lib/python3.9/site-packages/numpy/core/getlimits.py:549: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.\n",
      "  setattr(self, word, getattr(machar, word).flat[0])\n",
      "/home/lvyufeng/miniconda3/envs/mindspore/lib/python3.9/site-packages/numpy/core/getlimits.py:89: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.\n",
      "  return self._float_to_str(self.smallest_subnormal)\n",
      "Building prefix dict from the default dictionary ...\n",
      "Loading model from cache /tmp/jieba.cache\n",
      "Loading model cost 0.931 seconds.\n",
      "Prefix dict has been built successfully.\n",
      "/home/lvyufeng/miniconda3/envs/mindspore/lib/python3.9/site-packages/Cython/Compiler/Main.py:381: FutureWarning: Cython directive 'language_level' not set, using '3str' for now (Py3). This has changed from earlier releases! File: /home/lvyufeng/miniconda3/envs/mindspore/lib/python3.9/site-packages/mindnlp/transformers/models/graphormer/algos_graphormer.pyx\n",
      "  tree = Parsing.p_module(s, pxd, full_module_name)\n",
      "In file included from /home/lvyufeng/miniconda3/envs/mindspore/lib/python3.9/site-packages/numpy/core/include/numpy/ndarraytypes.h:1929,\n",
      "                 from /home/lvyufeng/miniconda3/envs/mindspore/lib/python3.9/site-packages/numpy/core/include/numpy/ndarrayobject.h:12,\n",
      "                 from /home/lvyufeng/miniconda3/envs/mindspore/lib/python3.9/site-packages/numpy/core/include/numpy/arrayobject.h:5,\n",
      "                 from /home/lvyufeng/.pyxbld/temp.linux-aarch64-cpython-39/home/lvyufeng/miniconda3/envs/mindspore/lib/python3.9/site-packages/mindnlp/transformers/models/graphormer/algos_graphormer.c:1240:\n",
      "/home/lvyufeng/miniconda3/envs/mindspore/lib/python3.9/site-packages/numpy/core/include/numpy/npy_1_7_deprecated_api.h:17:2: warning: #warning \"Using deprecated NumPy API, disable it with \" \"#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION\" [-Wcpp]\n",
      "   17 | #warning \"Using deprecated NumPy API, disable it with \" \\\n",
      "      |  ^~~~~~~\n"
     ]
    }
   ],
   "source": [
    "import mindspore\n",
    "from tqdm import tqdm\n",
    "from mindnlp import evaluate\n",
    "from mindnlp.dataset import load_dataset\n",
    "from mindnlp.transformers import AutoModelForSequenceClassification, AutoTokenizer\n",
    "from mindnlp.core.optim import AdamW\n",
    "from mindnlp.common.optimization import get_linear_schedule_with_warmup\n",
    "from mindnlp.peft import (\n",
    "    get_peft_model,\n",
    "    PeftType,\n",
    "    PromptTuningConfig,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e32c4a9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "model_name_or_path = \"roberta-large\"\n",
    "task = \"mrpc\"\n",
    "peft_type = PeftType.PROMPT_TUNING\n",
    "num_epochs = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "622fe9c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "peft_config = PromptTuningConfig(task_type=\"SEQ_CLS\", num_virtual_tokens=10)\n",
    "lr = 1e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "74e9efe0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/lvyufeng/miniconda3/envs/mindspore/lib/python3.9/site-packages/mindnlp/transformers/tokenization_utils_base.py:1526: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "if any(k in model_name_or_path for k in (\"gpt\", \"opt\", \"bloom\")):\n",
    "    padding_side = \"left\"\n",
    "else:\n",
    "    padding_side = \"right\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, padding_side=padding_side)\n",
    "if getattr(tokenizer, \"pad_token_id\") is None:\n",
    "    tokenizer.pad_token_id = tokenizer.eos_token_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "41a63e71-e7c4-4e5d-9e22-6953d981d4b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'sentence1': Tensor(shape=[], dtype=String, value= 'Amrozi accused his brother , whom he called \" the witness \" , of deliberately distorting his evidence .'), 'sentence2': Tensor(shape=[], dtype=String, value= 'Referring to him as only \" the witness \" , Amrozi accused his brother of deliberately distorting his evidence .'), 'label': Tensor(shape=[], dtype=Int64, value= 1), 'idx': Tensor(shape=[], dtype=Int64, value= 0)}\n"
     ]
    }
   ],
   "source": [
    "datasets = load_dataset(\"glue\", task)\n",
    "print(next(datasets['train'].create_dict_iterator()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bd2d7cd5-62b8-4b7a-ac69-338e6319152e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindnlp.dataset import BaseMapFunction\n",
    "\n",
    "class MapFunc(BaseMapFunction):\n",
    "    def __call__(self, sentence1, sentence2, label, idx):\n",
    "        outputs = tokenizer(sentence1, sentence2, truncation=True, max_length=None)\n",
    "        return outputs['input_ids'], outputs['attention_mask'], label\n",
    "\n",
    "\n",
    "def get_dataset(dataset, tokenizer):\n",
    "    input_colums=['sentence1', 'sentence2', 'label', 'idx']\n",
    "    output_columns=['input_ids', 'attention_mask', 'labels']\n",
    "    dataset = dataset.map(MapFunc(input_colums, output_columns),\n",
    "                          input_colums, output_columns)\n",
    "    dataset = dataset.padded_batch(batch_size, pad_info={'input_ids': (None, tokenizer.pad_token_id),\n",
    "                                                         'attention_mask': (None, 0)})\n",
    "    return dataset\n",
    "\n",
    "train_dataset = get_dataset(datasets['train'], tokenizer)\n",
    "eval_dataset = get_dataset(datasets['validation'], tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3b1fd5fc-2285-409e-a4e5-cc3c9759d77a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'input_ids': Tensor(shape=[32, 70], dtype=Int64, value=\n",
      "[[    0, 10127,  1001 ...     1,     1,     1],\n",
      " [    0,   975, 26802 ...     1,     1,     1],\n",
      " [    0,  1213,    56 ...     1,     1,     1],\n",
      " ...\n",
      " [    0,   133,  1154 ...     1,     1,     1],\n",
      " [    0, 12667,  8423 ...     1,     1,     1],\n",
      " [    0, 32478,  1033 ...     1,     1,     1]]), 'attention_mask': Tensor(shape=[32, 70], dtype=Int64, value=\n",
      "[[1, 1, 1 ... 0, 0, 0],\n",
      " [1, 1, 1 ... 0, 0, 0],\n",
      " [1, 1, 1 ... 0, 0, 0],\n",
      " ...\n",
      " [1, 1, 1 ... 0, 0, 0],\n",
      " [1, 1, 1 ... 0, 0, 0],\n",
      " [1, 1, 1 ... 0, 0, 0]]), 'labels': Tensor(shape=[32], dtype=Int64, value= [1, 0, 1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 1, 1, 1, 0, 1, \n",
      " 1, 1, 0, 0, 1, 1, 1, 0])}\n"
     ]
    }
   ],
   "source": [
    "print(next(train_dataset.create_dict_iterator()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "efb606a2-1fb5-415c-bf12-7e6fd324fe0a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "metric = evaluate.load(\"glue\", task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a3c15af0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MS_ALLOC_CONF]Runtime config:  enable_vmm:True  vmm_align_size:2MB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of RobertaForSequenceClassification were not initialized from the model checkpoint at roberta-large and are newly initialized: ['classifier.dense.bias', 'classifier.dense.weight', 'classifier.out_proj.bias', 'classifier.out_proj.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 1,061,890 || all params: 356,423,684 || trainable%: 0.2979291353713745\n"
     ]
    }
   ],
   "source": [
    "model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path, return_dict=True)\n",
    "model = get_peft_model(model, peft_config)\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6d3c5edb",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = AdamW(params=model.parameters(), lr=lr)\n",
    "\n",
    "# Instantiate scheduler\n",
    "lr_scheduler = get_linear_schedule_with_warmup(\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=0.06 * (len(train_dataset) * num_epochs),\n",
    "    num_training_steps=(len(train_dataset) * num_epochs),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d279225",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                                | 0/115 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [02:10<00:00,  1.13s/it]\n",
      " 15%|█████████████▋                                                                           | 2/13 [00:01<00:04,  2.21it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:02<00:00,  5.53it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 0: {'accuracy': 0.7205882352941176, 'f1': 0.8267477203647416}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:21<00:00,  1.41it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:02<00:00,  5.96it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 1: {'accuracy': 0.7009803921568627, 'f1': 0.817910447761194}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:22<00:00,  1.39it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:02<00:00,  5.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 2: {'accuracy': 0.7058823529411765, 'f1': 0.8198198198198198}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:20<00:00,  1.43it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:02<00:00,  5.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 3: {'accuracy': 0.7058823529411765, 'f1': 0.8187311178247734}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:23<00:00,  1.37it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:02<00:00,  5.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 4: {'accuracy': 0.7107843137254902, 'f1': 0.8190184049079755}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:23<00:00,  1.37it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:02<00:00,  5.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 5: {'accuracy': 0.7205882352941176, 'f1': 0.8161290322580645}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:23<00:00,  1.38it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:02<00:00,  5.80it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 6: {'accuracy': 0.7401960784313726, 'f1': 0.8295819935691319}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:24<00:00,  1.37it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:02<00:00,  5.71it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 7: {'accuracy': 0.7156862745098039, 'f1': 0.8104575163398693}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:22<00:00,  1.40it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:02<00:00,  5.89it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 8: {'accuracy': 0.7205882352941176, 'f1': 0.8093645484949833}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:15<00:00,  1.52it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:01<00:00,  7.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 9: {'accuracy': 0.7328431372549019, 'f1': 0.8256}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:09<00:00,  1.67it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:01<00:00,  6.92it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 10: {'accuracy': 0.7328431372549019, 'f1': 0.8233387358184765}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:05<00:00,  1.76it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:01<00:00,  6.91it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 11: {'accuracy': 0.7401960784313726, 'f1': 0.8284789644012945}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:05<00:00,  1.74it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:01<00:00,  7.02it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 12: {'accuracy': 0.7303921568627451, 'f1': 0.8264984227129337}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:02<00:00,  1.83it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:01<00:00,  6.91it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 13: {'accuracy': 0.7352941176470589, 'f1': 0.8296529968454258}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:04<00:00,  1.78it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:01<00:00,  6.79it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 14: {'accuracy': 0.7328431372549019, 'f1': 0.8244766505636071}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:06<00:00,  1.74it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:01<00:00,  6.59it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 15: {'accuracy': 0.7377450980392157, 'f1': 0.8260162601626017}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:08<00:00,  1.69it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:01<00:00,  6.58it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 16: {'accuracy': 0.7377450980392157, 'f1': 0.8288}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:07<00:00,  1.70it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:02<00:00,  6.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 17: {'accuracy': 0.7401960784313726, 'f1': 0.8295819935691319}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:08<00:00,  1.67it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:01<00:00,  6.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 18: {'accuracy': 0.7279411764705882, 'f1': 0.8235294117647058}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 115/115 [01:07<00:00,  1.71it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:01<00:00,  6.55it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 19: {'accuracy': 0.7328431372549019, 'f1': 0.8244766505636071}\n"
     ]
    }
   ],
   "source": [
    "from mindnlp.core import value_and_grad\n",
    "def forward_fn(**batch):\n",
    "    outputs = model(**batch)\n",
    "    loss = outputs.loss\n",
    "    return loss\n",
    "\n",
    "grad_fn = value_and_grad(forward_fn, tuple(model.parameters()))\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    model.set_train()\n",
    "    train_total_size = train_dataset.get_dataset_size()\n",
    "    for step, batch in enumerate(tqdm(train_dataset.create_dict_iterator(), total=train_total_size)):\n",
    "        optimizer.zero_grad()\n",
    "        loss = grad_fn(**batch)\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "\n",
    "    model.set_train(False)\n",
    "    eval_total_size = eval_dataset.get_dataset_size()\n",
    "    for step, batch in enumerate(tqdm(eval_dataset.create_dict_iterator(), total=eval_total_size)):\n",
    "        outputs = model(**batch)\n",
    "        predictions = outputs.logits.argmax(axis=-1)\n",
    "        predictions, references = predictions, batch[\"labels\"]\n",
    "        metric.add_batch(\n",
    "            predictions=predictions,\n",
    "            references=references,\n",
    "        )\n",
    "\n",
    "    eval_metric = metric.compute()\n",
    "    print(f\"epoch {epoch}:\", eval_metric)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f4ca2c1-d7a0-433d-8f9e-4436e4786a1a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  },
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
